import xarray as xr
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
from cartopy.io.shapereader import Reader
import numpy as np
import cartopy.feature as cfeat
import netCDF4 as nc
import matplotlib.colors as mcolors
from matplotlib.ticker import FuncFormatter
import os
from datetime import datetime,timedelta
from downloadgribs import downloadselect

#定义颜色映射
cmap_data = {
    'red':   [(0.0,  1.0, 1.0),
               # 白色开始(0.03,  1.0,1.0), 
              (0.2,  0.0, 0.0),
                  # 中间由白变蓝
              (0.7,  1.0, 1.0),
              (1.0,  0.0, 0.0)],  # 红色结束

    'green': [(0.0,  1.0, 1.0), 
               # 白色开始
              (0.2,  0.0, 0.0),
                (0.7,  0.0, 0.0),  # 中间过渡到蓝色此处为零
              (1.0,  0.0, 0.0)],  # 红色结束

    'blue':  [(0.0,  1.0, 1.0),
               # 白色开始
              (0.2,  1.0, 1.0),
               (0.7,  0.0, 0.0) ,  # 蓝色保持不变
              (1.0,  0.0, 0.0)]   # 红色结束
}


cmap = mcolors.LinearSegmentedColormap('custom', cmap_data)


plt.rcParams['font.sans-serif'] = ['SimHei'] # 用来正常显示中文标签SimHei
plt.rcParams['axes.unicode_minus'] = False # 用来正常显示负号


def grib2nc(gribpath,outncpath):
    

    gribname=os.path.basename(gribpath)   
    ncpath=os.path.join(outncpath,gribname+".nc")
    backend_kwargs={'filter_by_keys': {'typeOfLevel': 'surface','stepType': 'avg'}} 
    
    rh = xr.open_dataset(gribpath, engine='cfgrib',backend_kwargs=backend_kwargs)#,backend_kwargs=backend_kwargs
    print(rh["csnow"])
    rh["csnow"].to_netcdf(ncpath)




def createimages(band2,long,lati,outimagespath,foretime,starthours,endhours,extent1=[100,125, 17, 42]):
    extent=[extent1[0],extent1[2],extent1[1],extent1[3]]
    shppath=r'C:\基础数据\2019全国行政区划\省界.shp'
    shppath1=r'C:\基础数据\2019全国行政区划\市.shp'
    
    # --创建画图空间
    proj = ccrs.PlateCarree()  # 创建坐标系
    fig = plt.figure(figsize=(700,500),dpi=5) # 创建页面
    # plt.subplots_adjust(top=0.85)
    ax = fig.subplots(1, 1, subplot_kw={'projection': proj})  # 创建子图
    custom_levels = np.insert(np.arange(0,61,1).astype(np.float16),1,0.1)
    custom_levels2 = np.insert(np.arange(4,61,4).astype(np.float16),1,0.1)
 
    cf=ax.contourf(long, lati, band2,levels=custom_levels,cmap=cmap,extend='max')
    c3=fig.colorbar(cf,ax=ax,ticks=custom_levels2)
    c3.ax.tick_params(labelsize=500)
    c3.set_label("单位:mm",fontsize=800)
    # --设置地图属性
    reader = Reader(shppath)
    reader1 = Reader(shppath1)
    provinces = cfeat.ShapelyFeature(reader.geometries(), proj, edgecolor='k', facecolor='none')
    city = cfeat.ShapelyFeature(reader1.geometries(), proj, edgecolor='gray', facecolor='none')
    ax.add_feature(provinces, linewidth=40)
    ax.add_feature(city, linewidth=15)
    ax.set_extent(extent, crs=proj)

    ax.plot(115.82, 32.90, marker='o', color='black', markersize=250, alpha=0.7, transform=ccrs.Geodetic())
    ax.text(115.82-0.5, 32.90 + 0.5, "阜阳", verticalalignment='top', horizontalalignment='left',transform=ccrs.Geodetic(), color='black', fontsize=500)
    ax.plot(114.91, 25.86, marker='o', color='black', markersize=250, alpha=0.7, transform=ccrs.Geodetic())
    ax.text(114.91-0.5, 25.86 + 0.5, "赣州", verticalalignment='top', horizontalalignment='left',transform=ccrs.Geodetic(), color='black', fontsize=500)
    ax.plot(109.99, 27.61, marker='o', color='black', markersize=250, alpha=0.7, transform=ccrs.Geodetic())
    ax.text(109.99-0.5, 27.61 + 0.5, "怀化", verticalalignment='top', horizontalalignment='left',transform=ccrs.Geodetic(), color='black', fontsize=500)
    ax.plot(106.59, 23.89, marker='o', color='black', markersize=250, alpha=0.7, transform=ccrs.Geodetic())
    ax.text(106.59-0.5, 23.89 + 0.5, "百色", verticalalignment='top', horizontalalignment='left',transform=ccrs.Geodetic(), color='black', fontsize=500)
    ax.plot(121.42, 28.67, marker='o', color='black', markersize=250, alpha=0.7, transform=ccrs.Geodetic())
    ax.text(121.42-0.5, 28.67 + 0.5, "台州", verticalalignment='top', horizontalalignment='left',transform=ccrs.Geodetic(), color='black', fontsize=500)
    ax.plot(110.93, 21.67, marker='o', color='black', markersize=250, alpha=0.7, transform=ccrs.Geodetic())
    ax.text(110.93-0.5, 21.67 + 0.5, "茂名", verticalalignment='top', horizontalalignment='left',transform=ccrs.Geodetic(), color='black', fontsize=500)
    ax.plot(111.21, 34.77, marker='o', color='black', markersize=250, alpha=0.7, transform=ccrs.Geodetic())
    ax.text(111.21-0.5, 34.77 + 0.5, "三门峡", verticalalignment='top', horizontalalignment='left',transform=ccrs.Geodetic(), color='black', fontsize=500)
    ax.plot(118.68, 37.46, marker='o', color='black', markersize=250, alpha=0.7, transform=ccrs.Geodetic())
    ax.text(118.68-0.5, 37.46 + 0.5, "东营", verticalalignment='top', horizontalalignment='left',transform=ccrs.Geodetic(), color='black', fontsize=500)
    ax.plot(109.74, 38.31, marker='o', color='black', markersize=250, alpha=0.7, transform=ccrs.Geodetic())
    ax.text(109.74-0.5, 38.31 + 0.5, "榆林", verticalalignment='top', horizontalalignment='left',transform=ccrs.Geodetic(), color='black', fontsize=500)
    ax.plot(113.59, 37.89, marker='o', color='black', markersize=250, alpha=0.7, transform=ccrs.Geodetic())
    ax.text(113.59-0.5, 37.89 + 0.5, "阳泉", verticalalignment='top', horizontalalignment='left',transform=ccrs.Geodetic(), color='black', fontsize=500)
    ax.plot(114.33, 29.85, marker='o', color='black', markersize=250, alpha=0.7, transform=ccrs.Geodetic())
    ax.text(114.33-0.5, 29.85 + 0.5, "咸宁", verticalalignment='top', horizontalalignment='left',transform=ccrs.Geodetic(), color='black', fontsize=500)
    ax.plot(117.67, 24.52, marker='o', color='black', markersize=250, alpha=0.7, transform=ccrs.Geodetic())
    ax.text(117.67-0.5, 24.52 + 0.5, "漳州", verticalalignment='top', horizontalalignment='left',transform=ccrs.Geodetic(), color='black', fontsize=500)
    ax.plot(120.17, 33.35, marker='o', color='black', markersize=250, alpha=0.7, transform=ccrs.Geodetic())
    ax.text(120.17-0.5, 33.35 + 0.5, "盐城", verticalalignment='top', horizontalalignment='left',transform=ccrs.Geodetic(), color='black', fontsize=500)
    ax.plot(99.16, 25.17, marker='o', color='black', markersize=250, alpha=0.7, transform=ccrs.Geodetic())
    ax.text(99.16-0.5, 25.17 + 0.5, "保山", verticalalignment='top', horizontalalignment='left',transform=ccrs.Geodetic(), color='black', fontsize=500)
    ax.plot(105.58, 30.55, marker='o', color='black', markersize=250, alpha=0.7, transform=ccrs.Geodetic())
    ax.text(105.58-0.5, 30.55 + 0.5, "遂宁", verticalalignment='top', horizontalalignment='left',transform=ccrs.Geodetic(), color='black', fontsize=500)
    ax.plot(105.20, 37.52, marker='o', color='black', markersize=250, alpha=0.7, transform=ccrs.Geodetic())
    ax.text(105.20-0.5, 37.52 + 0.5, "中卫", verticalalignment='top', horizontalalignment='left',transform=ccrs.Geodetic(), color='black', fontsize=500)
    ax.plot(102.62, 37.95, marker='o', color='black', markersize=250, alpha=0.7, transform=ccrs.Geodetic())
    ax.text(102.62-0.5, 37.95 + 0.5, "武威", verticalalignment='top', horizontalalignment='left',transform=ccrs.Geodetic(), color='black', fontsize=500)
    ax.plot(110.39, 18.81, marker='o', color='black', markersize=250, alpha=0.7, transform=ccrs.Geodetic())
    ax.text(110.39-0.5, 18.81 + 0.5, "万宁", verticalalignment='top', horizontalalignment='left',transform=ccrs.Geodetic(), color='black', fontsize=500)
    ax.plot(120.20, 23.00, marker='o', color='black', markersize=250, alpha=0.7, transform=ccrs.Geodetic())
    ax.text(120.20-0.5, 23.00 + 0.5, "台南", verticalalignment='top', horizontalalignment='left',transform=ccrs.Geodetic(), color='black', fontsize=500)
    ax.plot(119.60, 39.96, marker='o', color='black', markersize=250, alpha=0.7, transform=ccrs.Geodetic())
    ax.text(119.60-0.5, 39.96 + 0.5, "秦皇岛", verticalalignment='top', horizontalalignment='left',transform=ccrs.Geodetic(), color='black', fontsize=500)


    ax.set_xticks(np.linspace(extent[0]+2, extent[1],5), crs=proj) 
    ax.tick_params(axis='x', labelsize=600)
    ax.set_yticks(np.linspace(extent[2]+2, extent[3],5),crs=proj)
    ax.tick_params(axis='y', labelsize=600)
#     def longitude_formatter(x, pos):
#         return f"{x}°E" if x >= 0 else f"{-x:.2f}°W"

#     def latitude_formatter(y, pos):
#         return f"{y}°N" if y >= 0 else f"{-y:.2f}°S"

# # 使用自定义的格式化函数
#     ax.xaxis.set_major_formatter(FuncFormatter(longitude_formatter))
#     ax.yaxis.set_major_formatter(FuncFormatter(latitude_formatter))
    
    starthours_str=str(starthours).rjust(3,'0')
    endhours_str=str(endhours).rjust(3,'0')
    ax.set_title(f"GFS预报模式{starthours_str}-{endhours_str}时累计降雪量",fontdict=dict(fontsize=1100, color='black'))
    fig.text(x=0.13,y=0.86,s="Total Snowfall SLR(10:1)",fontdict=dict(fontsize=550, color='black'))
    fig.text(x=0.63,y=0.86,s="Region:North China",fontdict=dict(fontsize=550, color='black'))#Central
    startforetime=datetime.strftime(datetime.strptime(foretime,"%Y%m%d%H"),"%Y%m%d%H")
    fig.text(x=0.1,y=0.09,s=f"起报时:{startforetime}(UTC)",fontdict=dict(fontsize=800, color='r'))

    startday=datetime.strftime(datetime.strptime(foretime,"%Y%m%d%H")+timedelta(hours=float(starthours)),"%Y%m%d%H")
    endday=datetime.strftime(datetime.strptime(foretime,"%Y%m%d%H")+timedelta(hours=float(endhours)),"%Y%m%d%H")
    fig.text(x=0.1,y=0.07,s=f"预报时:{startday}-{endday}(UTC)",fontdict=dict(fontsize=800, color='r'))
    fig.text(x=0.37,y=0.08,s="Program:雪落自序",fontdict=dict(fontsize=895, color='black'))
    fig.text(x=0.65,y=0.09,s="Model:GFS",fontdict=dict(fontsize=900, color='black'))
    fig.text(x=0.65,y=0.07,s="Producer:CCC",fontdict=dict(fontsize=900, color='black'))
    # plt.show()
    plt.savefig(outimagespath,bbox_inches='tight')


def grib2nc(forectime:str,starthours:int,endhours:int,outpath_prate:str, outpath_csnow:str):
    """grib转nc

    Args:
        forecdate (str): 预报起始时间
        starthours (int): 开始小时数
        endhours (int): 结束小时数
        outpath_prate (str): 降水量grib数据文件夹位置
        outpath_csnow (str): 雪分类grib数据文件夹位置
    """

    i=starthours
    while i<= endhours:

        hours=str(i).rjust(3,"0")

        fn = os.path.join(outpath_prate,f"gfs.t{forectime}z.pgrb2.0p25.f{hours}")
        backend_kwargs={'filter_by_keys': {'typeOfLevel': 'surface','stepType': 'avg'}} 
        #backend_kwargs={'filter_by_keys': {'typeOfLevel': 'surface', "cfVarName": "sde"}} 
        rh = xr.open_dataset(fn, engine='cfgrib',backend_kwargs=backend_kwargs)#,backend_kwargs=backend_kwargs
        print(rh['prate'])
        rh["prate"].to_netcdf(os.path.join(outpath_prate,f"gfs.t{forectime}z.pgrb2.0p25.f{hours}.nc"))
    
        fn = os.path.join(outpath_csnow,f"gfs.t{forectime}z.pgrb2.0p25.f{hours}")
        backend_kwargs={'filter_by_keys': {'typeOfLevel': 'surface','stepType': 'avg'}} 
        #backend_kwargs={'filter_by_keys': {'typeOfLevel': 'surface', "cfVarName": "sde"}} 
        rh = xr.open_dataset(fn, engine='cfgrib',backend_kwargs=backend_kwargs)#backend_kwargs=backend_kwargs
        print(rh["csnow"])
       
        rh["csnow"].to_netcdf(os.path.join(outpath_csnow,f"gfs.t{forectime}z.pgrb2.0p25.f{hours}.nc"))
        
        if i>=120:
            i+=3
        else:
            i+=1
    

def nc2png(forecdate:str,starthours:int,endhours:int,geoexten:tuple,outpath_prate:str, outpath_csnow:str,pngpath:str):
    """nc转png

    Args:
        forecdate (str): 预报起始时间
        starthours (int): 开始小时数
        endhours (int): 结束小时数
        outpath_prate (str): 降水量nc数据文件夹位置
        outpath_csnow (str): 雪分类nc数据文件夹位置
        pngpath (str): 输出png图片文件夹位置
    """

    forectime=forecdate[-2:]
    bands1=None
    lati=None
    long=None

    i=starthours
    while i<= endhours:
        hours=str(i).rjust(3,"0")
        fnpre = os.path.join(outpath_prate,f"gfs.t{forectime}z.pgrb2.0p25.f{hours}.nc")
        fncsonw = os.path.join(outpath_csnow,f"gfs.t{forectime}z.pgrb2.0p25.f{hours}.nc")
        nc_data_obj1 = nc.Dataset(fncsonw)
        band1 = np.asarray(nc_data_obj1.variables["csnow"])
        band1[band1!=1]=0
        
        print(band1)
    
        nc_data_obj1 = nc.Dataset(fnpre)
        if i > 120:
            band3 = np.asarray(nc_data_obj1.variables["prate"])*3600*3
        else:
            band3 = np.asarray(nc_data_obj1.variables["prate"])*3600

        band2=(band3*band1)
        if bands1 is None:
            bands1=np.zeros_like(band2)
        bands1 = bands1+band2
        long = nc_data_obj1.variables['longitude'][:]
        lati = nc_data_obj1.variables['latitude'][:]
    
        if i>=120:
            i+=3
        else:
            i+=1

    createimages(bands1,long,lati,os.path.join(pngpath,f"allsnow{str(starthours)}_{str(endhours)}.png"),forecdate,starthours,endhours,geoexten)

if __name__=="__main__":
    

    downloadselect("2024020500",240,384, [95,17,125,42], r"C:\Users\DELL\Desktop\pre", r"C:\Users\DELL\Desktop\csnow")#[95,125,17,42]
    grib2nc("00",240,384,r"C:\Users\DELL\Desktop\pre",r"C:\Users\DELL\Desktop\csnow")
    nc2png("2024020500",240,384,[95,17,125,42],r"C:\Users\DELL\Desktop\pre",r"C:\Users\DELL\Desktop\csnow",r"C:\Users\DELL\Desktop")
   